package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"os"
	"strings"
)

// 定义命令行参数
var (
	inputFile  = flag.String("input-file", "", "输入文件路径")
	outputFile = flag.String("output-file", "", "输出文件路径")
	batchSize  = flag.Int("batch", 100, "批量合并的行数")
	mode       = flag.String("mode", "merge", "操作模式: merge(合并) 或 split(拆分)")
	title      = flag.String("title", "", "合并时的语句前缀类型: insert/ignore/replace")
)

type SQLStatement struct {
	Type   string   // INSERT, INSERT IGNORE, REPLACE
	Table  string   // 表名
	Fields string   // 字段列表
	Values []string // 值列表
	Line   string   // 原始行内容（用于保持原样）
}

func main() {
	flag.Parse()

	// 检查必需的参数
	if *inputFile == "" || *outputFile == "" {
		fmt.Println("错误: 必须指定 --input-file 和 --output-file 参数")
		flag.PrintDefaults()
		os.Exit(1)
	}

	// 检查 mode 参数是否有效
	if *mode != "merge" && *mode != "split" {
		fmt.Println("错误: mode 参数必须是 merge 或 split")
		os.Exit(1)
	}
	
	// 检查 title 参数是否有效
	if *title != "" && *title != "insert" && *title != "ignore" && *title != "replace" {
		fmt.Println("错误: title 参数必须是 insert, ignore 或 replace")
		os.Exit(1)
	}

	// 打开输入文件
	input, err := os.Open(*inputFile)
	if err != nil {
		fmt.Printf("无法打开输入文件: %v\n", err)
		os.Exit(1)
	}
	defer input.Close()

	// 创建输出文件
	output, err := os.Create(*outputFile)
	if err != nil {
		fmt.Printf("无法创建输出文件: %v\n", err)
		os.Exit(1)
	}
	defer output.Close()

	if *mode == "merge" {
		err = mergeStatements(input, output)
	} else {
		err = splitStatements(input, output)
	}

	if err != nil {
		fmt.Printf("处理文件时出错: %v\n", err)
		os.Exit(1)
	}

	fmt.Println("文件处理完成")
}

// mergeStatements 将单条INSERT语句合并为批量语句
func mergeStatements(input *os.File, output *os.File) error {
	reader := bufio.NewReader(input)
	writer := bufio.NewWriter(output)
	defer writer.Flush()

	var currentTable string
	var currentType string
	var currentFields string
	var linesToMerge []string // 存储需要合并的原始行

	for {
		line, err := reader.ReadString('\n')
		if err != nil && err != io.EOF {
			return err
		}

		// 检查是否是INSERT语句
		if isInsertStatement(line) {
			// 检查是否是INSERT ... SELECT语句
			if isInsertSelectStatement(line) {
				// INSERT ... SELECT语句无法合并，直接写入之前的批量语句并写入当前行
				if len(linesToMerge) > 0 {
					mergeAndWriteLines(writer, linesToMerge, currentType, currentTable, currentFields)
					linesToMerge = nil
				}
				writer.WriteString(line)
				if err == io.EOF {
					break
				}
				continue
			}
			
			stmt := parseInsertStatement(line)
			if stmt == nil {
				// 如果不是有效的INSERT语句，写入之前的批量语句并写入当前行
				if len(linesToMerge) > 0 {
					mergeAndWriteLines(writer, linesToMerge, currentType, currentTable, currentFields)
					linesToMerge = nil
				}
				writer.WriteString(line)
				if err == io.EOF {
					break
				}
				continue
			}

			// 如果是相同表和类型的语句，收集行
			if stmt.Table == currentTable && stmt.Type == currentType && stmt.Fields == currentFields {
				linesToMerge = append(linesToMerge, line)
			} else {
				// 如果是不同表或类型的语句，先写入之前的批量语句
				if len(linesToMerge) > 0 {
					mergeAndWriteLines(writer, linesToMerge, currentType, currentTable, currentFields)
				}
				// 开始新的表和类型
				currentTable = stmt.Table
				currentType = stmt.Type
				currentFields = stmt.Fields
				linesToMerge = []string{line}
			}
		} else {
			// 不是INSERT语句，写入之前的批量语句
			if len(linesToMerge) > 0 {
				mergeAndWriteLines(writer, linesToMerge, currentType, currentTable, currentFields)
				linesToMerge = nil
				currentTable = ""
				currentType = ""
				currentFields = ""
			}
			// 写入非INSERT语句
			writer.WriteString(line)
		}

		if err == io.EOF {
			// 处理最后一批语句
			if len(linesToMerge) > 0 {
				mergeAndWriteLines(writer, linesToMerge, currentType, currentTable, currentFields)
			}
			break
		}
	}

	return nil
}

// splitStatements 将批量INSERT语句拆分为单条语句
func splitStatements(input *os.File, output *os.File) error {
	reader := bufio.NewReader(input)
	writer := bufio.NewWriter(output)
	defer writer.Flush()

	for {
		line, err := reader.ReadString('\n')
		if err != nil && err != io.EOF {
			return err
		}

		// 检查是否是批量INSERT语句
		if isBatchInsertStatement(line) {
			stmt := parseBatchInsertStatement(line)
			if stmt != nil {
				// 拆分为单条语句
				for _, value := range stmt.Values {
					writer.WriteString(fmt.Sprintf("%s INTO %s %s VALUES %s;\n", 
						stmt.Type, stmt.Table, stmt.Fields, value))
				}
			} else {
				// 如果解析失败，原样写入
				writer.WriteString(line)
			}
		} else {
			// 不是批量INSERT语句，原样写入
			writer.WriteString(line)
		}

		if err == io.EOF {
			break
		}
	}

	return nil
}

// isInsertStatement 检查是否是INSERT语句
func isInsertStatement(line string) bool {
	trimmed := strings.TrimSpace(line)
	lowerLine := strings.ToLower(trimmed)
	return strings.HasPrefix(lowerLine, "insert into") ||
		strings.HasPrefix(lowerLine, "insert ignore into") ||
		strings.HasPrefix(lowerLine, "replace into")
}

// isBatchInsertStatement 检查是否是批量INSERT语句
func isBatchInsertStatement(line string) bool {
	trimmed := strings.TrimSpace(line)
	if !isInsertStatement(trimmed) {
		return false
	}
	
	// 检查是否包含多个VALUES，或者单个VALUES中包含多个括号组
	lowerLine := strings.ToLower(trimmed)
	valuesIndex := strings.Index(lowerLine, "values")
	
	if valuesIndex == -1 {
		return false
	}
	
	// 提取VALUES后面的部分
	valuesPart := trimmed[valuesIndex+len("values"):]
	
	// 查找分号位置
	semiColonIndex := strings.Index(valuesPart, ";")
	if semiColonIndex != -1 {
		valuesPart = valuesPart[:semiColonIndex]
	}
	
	valuesPart = strings.TrimSpace(valuesPart)
	
	// 检查是否包含逗号和多个括号组，如 (val1), (val2)
	return strings.Contains(valuesPart, "),(") || 
		   (strings.Count(valuesPart, "(") > 1 && strings.Contains(valuesPart, ","))
}

// isInsertSelectStatement 检查是否是INSERT ... SELECT语句
func isInsertSelectStatement(line string) bool {
	trimmed := strings.TrimSpace(line)
	lowerLine := strings.ToLower(trimmed)
	
	// 检查是否包含SELECT关键字且在VALUES之后或没有VALUES
	selectIndex := strings.Index(lowerLine, "select")
	valuesIndex := strings.Index(lowerLine, "values")
	
	// 如果有SELECT且没有VALUES，或者SELECT在VALUES之后，则是INSERT ... SELECT语句
	return selectIndex != -1 && (valuesIndex == -1 || selectIndex > valuesIndex)
}

// parseInsertStatement 解析单条INSERT语句
func parseInsertStatement(line string) *SQLStatement {
	originalLine := line
	line = strings.TrimSpace(line)
	stmt := &SQLStatement{
		Line: originalLine,
	}

	// 提取INSERT语句的类型和原始格式
	var typePrefix string
	if strings.HasPrefix(strings.ToLower(line), "insert ignore into") {
		stmt.Type = "INSERT IGNORE"
		typePrefix = line[:len("insert ignore into")]
		line = line[len("insert ignore into"):]
	} else if strings.HasPrefix(strings.ToLower(line), "insert into") {
		stmt.Type = "INSERT"
		typePrefix = line[:len("insert into")]
		line = line[len("insert into"):]
	} else if strings.HasPrefix(strings.ToLower(line), "replace into") {
		stmt.Type = "REPLACE"
		typePrefix = line[:len("replace into")]
		line = line[len("replace into"):]
	} else {
		return nil
	}

	stmt.Line = typePrefix + line // 保存原始格式的完整语句

	line = strings.TrimSpace(line)
	
	// 检查是否是INSERT ... SELECT语句
	selectIndex := strings.Index(strings.ToLower(line), "select")
	valuesIndex := strings.Index(strings.ToLower(line), "values")
	
	// 如果有SELECT但没有VALUES，这是INSERT ... SELECT语句
	if selectIndex != -1 && (valuesIndex == -1 || selectIndex < valuesIndex) {
		// 这是INSERT ... SELECT语句
		stmt.Table = strings.TrimSpace(line[:selectIndex])
		stmt.Fields = ""
		return stmt
	}
	
	// 处理INSERT ... VALUES语句
	if valuesIndex == -1 {
		return nil
	}
	
	// 提取表名（VALUES之前的部分）
	tablePart := strings.TrimSpace(line[:valuesIndex])
	
	// 检查表部分是否包含括号（字段定义）
	tableEnd := strings.Index(tablePart, "(")
	if tableEnd == -1 {
		// 没有显式字段
		stmt.Table = tablePart
		stmt.Fields = ""
	} else {
		// 有显式字段
		stmt.Table = strings.TrimSpace(tablePart[:tableEnd])
		
		// 提取字段部分
		fieldsStart := tableEnd
		fieldsEnd := findMatchingParen(tablePart, fieldsStart)
		if fieldsEnd == -1 {
			// 如果找不到匹配的括号，整个表部分作为表名
			stmt.Table = tablePart
			stmt.Fields = ""
		} else {
			stmt.Fields = tablePart[fieldsStart : fieldsEnd+1]
		}
	}
	
	return stmt
}

// parseBatchInsertStatement 解析批量INSERT语句
func parseBatchInsertStatement(line string) *SQLStatement {
	originalLine := line
	line = strings.TrimSpace(line)
	stmt := &SQLStatement{
		Line: originalLine,
	}

	// 提取INSERT语句的类型
	if strings.HasPrefix(strings.ToLower(line), "insert ignore into") {
		stmt.Type = "INSERT IGNORE"
		line = line[len("insert ignore into"):]
	} else if strings.HasPrefix(strings.ToLower(line), "insert into") {
		stmt.Type = "INSERT"
		line = line[len("insert into"):]
	} else if strings.HasPrefix(strings.ToLower(line), "replace into") {
		stmt.Type = "REPLACE"
		line = line[len("replace into"):]
	} else {
		return nil
	}

	line = strings.TrimSpace(line)
	
	// 查找VALUES关键字
	valuesIndex := strings.Index(strings.ToLower(line), "values")
	if valuesIndex == -1 {
		return nil
	}
	
	// 提取表名（VALUES之前的部分）
	tablePart := strings.TrimSpace(line[:valuesIndex])
	
	// 检查表部分是否包含括号（字段定义）
	tableEnd := strings.Index(tablePart, "(")
	if tableEnd == -1 {
		// 没有显式字段
		stmt.Table = tablePart
		stmt.Fields = ""
	} else {
		// 有显式字段
		stmt.Table = strings.TrimSpace(tablePart[:tableEnd])
		
		// 提取字段部分
		fieldsStart := tableEnd
		fieldsEnd := findMatchingParen(tablePart, fieldsStart)
		if fieldsEnd == -1 {
			// 如果找不到匹配的括号，整个表部分作为表名
			stmt.Table = tablePart
			stmt.Fields = ""
		} else {
			stmt.Fields = tablePart[fieldsStart : fieldsEnd+1]
		}
	}
	
	// 提取VALUES部分
	valuesPart := line[valuesIndex+len("values"):]
	
	// 查找分号并移除
	semiColonIndex := strings.Index(valuesPart, ";")
	if semiColonIndex != -1 {
		valuesPart = valuesPart[:semiColonIndex]
	}
	
	valuesPart = strings.TrimSpace(valuesPart)
	
	// 解析值
	stmt.Values = parseValues(valuesPart)
	
	return stmt
}

// parseValues 解析VALUES部分
func parseValues(line string) []string {
	line = strings.TrimSpace(line)
	if len(line) == 0 {
		return nil
	}

	var result []string
	var current strings.Builder
	var parenCount int

	for _, char := range line {
		switch char {
		case '(':
			if parenCount == 0 {
				current.Reset()
			}
			parenCount++
			current.WriteRune(char)
		case ')':
			current.WriteRune(char)
			parenCount--
			if parenCount == 0 {
				result = append(result, current.String())
			}
		default:
			if parenCount > 0 {
				current.WriteRune(char)
			}
		}
	}

	return result
}

// findMatchingParen 查找匹配的右括号
func findMatchingParen(s string, start int) int {
	if start >= len(s) || s[start] != '(' {
		return -1
	}
	
	parenCount := 1
	for i := start + 1; i < len(s); i++ {
		if s[i] == '(' {
			parenCount++
		} else if s[i] == ')' {
			parenCount--
			if parenCount == 0 {
				return i
			}
		}
	}
	
	return -1
}

// mergeAndWriteLines 合并并写入行
func mergeAndWriteLines(writer *bufio.Writer, lines []string, stmtType, table, fields string) error {
	if len(lines) == 0 {
		return nil
	}

	// 如果只有一行，则直接输出原始行，不做任何修改
	if len(lines) == 1 {
		// 直接写入原始行，但确保以换行符结尾
		line := lines[0]
		if !strings.HasSuffix(line, "\n") {
			line += "\n"
		}
		_, err := writer.WriteString(line)
		return err
	}

	// 解析所有行的VALUES部分
	var allValues []string
	for _, line := range lines {
		values := extractValuesFromLine(line)
		allValues = append(allValues, values...)
	}

	// 确定输出的语句类型前缀
	outputType := stmtType
	if *title != "" {
		switch *title {
		case "insert":
			outputType = "INSERT"
		case "ignore":
			outputType = "INSERT IGNORE"
		case "replace":
			outputType = "REPLACE"
		}
	}

	// 按batchSize分批写入
	for i := 0; i < len(allValues); i += *batchSize {
		end := i + *batchSize
		if end > len(allValues) {
			end = len(allValues)
		}

		batchValues := allValues[i:end]

		// 写入语句头
		var err error
		if fields == "" {
			// 当没有字段时，不添加额外的空格
			_, err = fmt.Fprintf(writer, "%s INTO %s VALUES ", outputType, table)
		} else {
			// 当有字段时，添加字段
			_, err = fmt.Fprintf(writer, "%s INTO %s %s VALUES ", outputType, table, fields)
		}
		if err != nil {
			return err
		}

		// 写入VALUES
		for j, value := range batchValues {
			if j > 0 {
				_, err = writer.WriteString(", ")
				if err != nil {
					return err
				}
			}
			_, err = writer.WriteString(value)
			if err != nil {
				return err
			}
		}

		_, err = writer.WriteString(";\n")
		if err != nil {
			return err
		}
	}

	return nil
}

// extractValuesFromLine 从行中提取VALUES部分
func extractValuesFromLine(line string) []string {
	line = strings.TrimSpace(line)
	
	// 查找VALUES关键字
	valuesIndex := strings.Index(strings.ToLower(line), "values")
	if valuesIndex == -1 {
		return nil
	}
	
	// 提取VALUES部分
	valuesPart := line[valuesIndex+len("values"):]
	
	// 查找分号并移除
	semiColonIndex := strings.Index(valuesPart, ";")
	if semiColonIndex != -1 {
		valuesPart = valuesPart[:semiColonIndex]
	}
	
	valuesPart = strings.TrimSpace(valuesPart)
	
	// 解析值
	return parseValues(valuesPart)
}